﻿//------------------------------------------------------------------------------
// <copyright file="Program.cs" company="Microsoft">
//     Copyright (c) Microsoft Corporation.  All rights reserved.
// </copyright>
//------------------------------------------------------------------------------

namespace Rpc.Tests.Launcher
{    
    using System;
    using System.Diagnostics;
    using System.IO;
    using System.Linq;
    using System.Text;
    using Microsoft.Win32;

    /// <summary>
    /// Used to launch applications as administrator when registered as a scheduled task.
    /// UAC prompt is only displayed at the time of registering, not every time when launching.
    /// Because tasks cannot be passed a command line, it is specified through the registry.
    /// </summary>
    public class Program
    {
        private const string LaunchElevatedTaskName = "RpcTestsLaunchElevated";

        /// <summary>
        /// The application entry point.
        /// </summary>
        /// <param name="args">Command line arguments (unused).</param>
        public static void Main(string[] args)
        {
            if (args.Contains("/register"))
            {
                // If registering, create a scheduled task (this will display UAC prompt) and exit
                Register();
                return;
            }

            // Get the command line to execute from the registry
            using (RegistryKey key = Registry.CurrentUser.OpenSubKey(@"Software\Microsoft\MSTDKRPC.Tests"))
            {
                string commandLine = key.GetValue("LauncherCommandLine").ToString();

                if (!string.IsNullOrEmpty(commandLine) && commandLine.Length > 0)
                {
                    // Split into file name and arguments, accounting for quoted paths
                    int sep = 0;
                    bool quote = false;

                    while (sep < commandLine.Length)
                    {
                        if (commandLine[sep] == '\"')
                        {
                            quote = !quote;
                        }
                        else if (commandLine[sep] == ' ' && !quote)
                        {
                            break;
                        }

                        ++sep;
                    }

                    if (sep == commandLine.Length)
                    {
                        sep = -1;
                    }

                    string fileName = null, arguments = null;

                    if (sep != -1)
                    {
                        fileName = commandLine.Substring(0, sep).Replace("\"", string.Empty);
                        arguments = commandLine.Substring(sep + 1);
                    }
                    else
                    {
                        fileName = commandLine;
                    }

                    // Start the process using extracted file name and arguments
                    Process.Start(new ProcessStartInfo
                    {
                        FileName = fileName,
                        Arguments = arguments,
                        UseShellExecute = true,
                        WindowStyle = ProcessWindowStyle.Hidden
                    }).WaitForExit();
                }
            }
        }

        /// <summary>
        /// Registers a scheduled task to run the launcher if not already registered.
        /// This will bring up the UAC prompt.
        /// </summary>
        public static void Register()
        {
            // Unregister any previous instances
            Process.Start(new ProcessStartInfo
            {
                WindowStyle = ProcessWindowStyle.Hidden,
                UseShellExecute = true,
                FileName = "cmd.exe",
                Arguments = string.Format("/c schtasks.exe /delete /tn {0}", LaunchElevatedTaskName)
            }).WaitForExit();

            // Create a task with manual scheduling trigger using the launcher
            string launcherPath = System.Reflection.Assembly.GetExecutingAssembly().Location;
            string tempTaskFile = Path.Combine(Path.GetTempPath(), string.Format("{0}.xml", LaunchElevatedTaskName));
            string taskContents = string.Format(
                "<?xml version=\"1.0\" encoding=\"UTF-16\"?>" +
                    "<Task version=\"1.2\" xmlns=\"http://schemas.microsoft.com/windows/2004/02/mit/task\">" +
                        "<Triggers />" +
                        "<Principals>" +
                            "<Principal id=\"Author\">" +
                                "<UserId>{0}\\{1}</UserId>" +
                                "<LogonType>InteractiveToken</LogonType>" +
                                "<RunLevel>HighestAvailable</RunLevel>" +
                            "</Principal>" +
                        "</Principals>" +
                        "<Settings>" +
                            "<MultipleInstancesPolicy>Parallel</MultipleInstancesPolicy>" +
                            "<DisallowStartIfOnBatteries>false</DisallowStartIfOnBatteries>" +
                            "<StopIfGoingOnBatteries>false</StopIfGoingOnBatteries>" +
                            "<AllowHardTerminate>true</AllowHardTerminate>" +
                            "<StartWhenAvailable>false</StartWhenAvailable>" +
                            "<RunOnlyIfNetworkAvailable>false</RunOnlyIfNetworkAvailable>" +
                            "<IdleSettings>" +
                                "<StopOnIdleEnd>false</StopOnIdleEnd>" +
                                "<RestartOnIdle>false</RestartOnIdle>" +
                            "</IdleSettings>" +
                            "<AllowStartOnDemand>true</AllowStartOnDemand>" +
                            "<Enabled>true</Enabled>" +
                            "<Hidden>false</Hidden>" +
                            "<RunOnlyIfIdle>false</RunOnlyIfIdle>" +
                            "<WakeToRun>false</WakeToRun>" +
                            "<ExecutionTimeLimit>P3D</ExecutionTimeLimit>" +
                            "<Priority>4</Priority>" +
                      "</Settings>" +
                      "<Actions Context=\"Author\">" +
                      "<Exec>" +
                        "<Command>{2}</Command>" +
                      "</Exec>" +
                      "</Actions>" +
                      "</Task>",
                      Environment.UserDomainName,
                      Environment.UserName,
                      launcherPath);

            File.WriteAllText(tempTaskFile, taskContents, Encoding.Unicode);

            Process.Start(new ProcessStartInfo
            {
                Verb = "runas",
                WindowStyle = ProcessWindowStyle.Hidden,
                UseShellExecute = true,
                FileName = "cmd.exe",
                Arguments = string.Format("/c schtasks.exe /create /tn {0} /xml \"{1}\"", LaunchElevatedTaskName, tempTaskFile)
            }).WaitForExit();

            File.Delete(tempTaskFile);
        }
    }
}
